package org.finesys.chat.core.base.embedding.store.redis;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.finesys.chat.core.base.document.Metadata;
import org.finesys.chat.core.base.embedding.Embedding;
import org.finesys.chat.core.base.embedding.EmbeddingMatch;
import org.finesys.chat.core.base.embedding.EmbeddingSearchRequest;
import org.finesys.chat.core.base.embedding.EmbeddingSearchResult;
import org.finesys.chat.core.base.embedding.EmbeddingStore;
import org.finesys.chat.core.base.embedding.filter.Filter;
import org.finesys.chat.core.base.embedding.store.util.Generator;
import org.finesys.chat.core.base.segment.TextSegment;
import org.finesys.commom.jackson.util.JacksonUtil;
import org.finesys.common.core.util.Utils;
import org.finesys.common.core.util.ValidationUtil;

import lombok.extern.slf4j.Slf4j;
import redis.clients.jedis.JedisPooled;
import redis.clients.jedis.Pipeline;
import redis.clients.jedis.params.ScanParams;
import redis.clients.jedis.resps.ScanResult;
import redis.clients.jedis.search.Document;
import redis.clients.jedis.search.FTCreateParams;
import redis.clients.jedis.search.IndexDataType;
import redis.clients.jedis.search.Query;
import redis.clients.jedis.search.RediSearchUtil;
import redis.clients.jedis.search.SearchResult;
import redis.clients.jedis.search.schemafields.SchemaField;


@Slf4j
public class RedisEmbeddingStore implements EmbeddingStore<TextSegment>, AutoCloseable {

    private static final String QUERY_TEMPLATE = "%s=>[ KNN %d @%s $BLOB AS %s]";
    private final JedisPooled client;
    private final RedisSchema schema;
    private final RedisMetadataFilterMapper filterMapper;

    public RedisEmbeddingStore(String host, int port, String username, String password, String indexName, String prefix, Integer dimension, Map<String, SchemaField> metadataConfig) {
        ValidationUtil.ensureNotBlank(host, "host");
        ValidationUtil.ensureGreaterThanZero(port, "port");
        this.client = username == null ? new JedisPooled(host, port) : new JedisPooled(host, port, username, password);
        this.schema = RedisSchema.builder().indexName(indexName).prefix(prefix).dimension(dimension).metadataConfig(metadataConfig).build();
        this.filterMapper = new RedisMetadataFilterMapper(metadataConfig);
        if (!this.isIndexExist(this.schema.indexName())) {
            ValidationUtil.ensureNotNull(dimension, "dimension");
            this.createIndex(this.schema.indexName());
        }
    }

    public RedisEmbeddingStore(String uri, String indexName, String prefix, Integer dimension, Map<String, SchemaField> metadataConfig) {
        ValidationUtil.ensureNotBlank(uri, "uri");
        this.client = new JedisPooled(uri);
        this.schema = RedisSchema.builder().indexName((String) Utils.getOrDefault(indexName, "embedding-index")).prefix((String) Utils.getOrDefault(prefix, "embedding:")).dimension(dimension).metadataConfig(Utils.copyIfNotNull(metadataConfig)).build();
        this.filterMapper = new RedisMetadataFilterMapper(metadataConfig);
        if (!this.isIndexExist(this.schema.indexName())) {
            ValidationUtil.ensureNotNull(dimension, "dimension");
            this.createIndex(this.schema.indexName());
        }
    }

    private void createIndex(String indexName) {
        String res = this.client.ftCreate(indexName, FTCreateParams.createParams().on(IndexDataType.JSON)
                .addPrefix(this.schema.prefix()), this.schema.toSchemaFields());
        if (!"OK".equals(res)) {
            if (log.isErrorEnabled()) {
                log.error("create index error, msg={}", res);
            }
            throw new RedisRequestFailedException("create index error, msg=" + res);
        }
    }

    /**
     * 判断指定索引是否存在
     *
     * @param indexName 要检查的索引名称
     * @return 如果索引存在返回true，否则返回false
     */
    private boolean isIndexExist(String indexName) {
        try {
            this.client.ftInfo(indexName);
            return true;
        } catch (Exception e) {
            return false;
        }
    }


    @Override
    public String add(Embedding embedding) {
        String id = UUID.randomUUID().toString();
        return this.add(embedding, id);
    }

    @Override
    public String add(Embedding embedding, String key) {
        return this.addInternal(key, embedding, null);
    }


    @Override
    public String add(Embedding embedding, TextSegment textSegment) {
        String id = UUID.randomUUID().toString();
        return this.addInternal(id, embedding, textSegment);
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings, List<TextSegment> textSegments) {
        List<String> ids = Generator.generateRandomIds(embeddings.size());
        return this.addAll(ids, embeddings, textSegments);
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings) {
        List<String> ids = Generator.generateRandomIds(embeddings.size());
        return this.addAll(ids, embeddings, null);
    }

    @Override
    public List<String> addAll(List<String> ids, List<Embedding> embeddings, List<TextSegment> textSegments) {
        ValidationUtil.ensureNotEmpty(ids, "ids");
        ValidationUtil.ensureNotEmpty(embeddings, "embeddings");
        ValidationUtil.ensureTrue(ids.size() == embeddings.size(), "ids and embeddings size must be equal");
        ValidationUtil.ensureTrue(textSegments == null || textSegments.size() == embeddings.size(), "textSegments and embeddings size must be equal");
        Pipeline pipeline = this.client.pipelined();

        List<Object> responses = new ArrayList<>();
        try {
            for (int i = 0; i < ids.size(); i++) {
                String id = ids.get(i);
                Embedding embedding = embeddings.get(i);
                TextSegment textSegment = textSegments == null ? null : textSegments.get(i);
                Map<String, Object> fields = new HashMap<>();
                fields.put(this.schema.vectorFieldName(), embedding.getVector());
                if (textSegment != null) {
                    fields.put(this.schema.scalarFieldName(), textSegment.getText());
                    fields.putAll(textSegment.getMetadata().toMap());
                }
                String key = String.format("%s%s", this.schema.prefix(), id);
                pipeline.jsonSetWithEscape(key, RedisSchema.JSON_SET_PATH, fields);
                //最后一个
                if (i == ids.size() - 1) {
                    responses = pipeline.syncAndReturnAll();
                }
                i++;
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            if (pipeline != null) {
                pipeline.close();
            }
        }
        Optional<Object> errResponse = responses.stream().filter(response -> !"OK".equals(response)).findAny();
        if (errResponse.isPresent()) {
            if (log.isErrorEnabled()) {
                log.error("add embedding failed, msg={}", errResponse.get());
            }
            throw new RedisRequestFailedException("add embedding failed, msg=" + String.valueOf(errResponse.get()));
        }
        return ids;
    }

    private String addInternal(String key, Embedding embedding, TextSegment textSegment) {
        this.addAll(Collections.singletonList(key), Collections.singletonList(embedding), textSegment == null ? null : Collections.singletonList(textSegment));
        return key;
    }

    @Override
    public void delete(String key) {
        ValidationUtil.ensureNotBlank(key, "key");
        this.deleteAll(Collections.singletonList(key));
    }

    @Override
    public void deleteAll(List<String> keys) {
        ValidationUtil.ensureNotEmpty(keys, "keys");
        String[] redisKeys = keys.stream().map(k -> String.format("%s%s", this.schema.prefix(), k)).toArray(String[]::new);
        this.client.del(redisKeys);
    }

    @Override
    public void deleteAll(Filter filter) {
        ValidationUtil.ensureNotNull(filter, "filter");
        SearchResult result = this.client.ftSearch(this.schema.indexName(), this.filterMapper.map(filter));
        String[] keys = result.getDocuments().stream().map(Document::getId).toArray(String[]::new);
        this.client.del(keys);
    }

    @Override
    public void deleteAll() {
        //避免直接删除所有key，因为redis集群模式下，会导致其他节点数据丢失
        ScanParams scanParams = new ScanParams().match(this.schema.prefix() + "*").count(1000);
        String nextCursor = "0";
        do {
            //scan 命令返回的是一个游标，下一次扫描需要用到上一次的游标作为参数
            ScanResult<String> scanResult = this.client.scan(nextCursor, scanParams);
            List<String> keysList = scanResult.getResult();
            if (!keysList.isEmpty()) {
                this.client.del(keysList.toArray(new String[0]));
            }
            nextCursor = scanResult.getCursor();
        } while (!"0".equals(nextCursor));
    }

    @Override
    public EmbeddingSearchResult<TextSegment> search(EmbeddingSearchRequest request) {
        Query query = new Query(String.format(QUERY_TEMPLATE, this.filterMapper.map(request.getFilter()), request.getLimit(), this.schema.vectorFieldName(), "vector_score"))
                .addParam("BLOB", RediSearchUtil.toByteArray(request.getQueryEmbedding().getVectorFloat()))
                .setSortBy("vector_score", true)
                .limit(0, request.getLimit())
                .dialect(2);
        SearchResult result = this.client.ftSearch(this.schema.indexName(), query);
        List<Document> documents = result.getDocuments();
        return new EmbeddingSearchResult<>(this.toEmbeddingMatch(documents, request.getMinScore()));
    }

    private List<EmbeddingMatch<TextSegment>> toEmbeddingMatch(List<Document> documents, double minScore) {
        List<EmbeddingMatch<TextSegment>> matches = documents.stream().map(doc -> {
            String id = doc.getId();
            double score = (2.0 - Double.parseDouble(doc.getString("vector_score"))) / 2.0;
            Map<String, Object> properties = JacksonUtil.parseMap(doc.getString("$"));
            ValidationUtil.ensureNotNull(properties, "properties");
            List<Double> vectors = (List<Double>) properties.get(this.schema.vectorFieldName());
            Embedding embedding = Embedding.from(vectors);
            String text = properties.containsKey(this.schema.scalarFieldName()) ? properties.get(this.schema.scalarFieldName()).toString() : null;
            TextSegment textSegment = null;
            if (text != null) {
                Map<String, Object> metadata = this.schema.metadataConfig().keySet().stream().filter(properties::containsKey)
                        .collect(Collectors.toMap(Function.identity(), properties::get));
                textSegment = TextSegment.from(text, Metadata.from(metadata));
            }
            return new EmbeddingMatch<>(score, id, embedding, textSegment, 0);
        }).collect(Collectors.toList());
        return matches;
    }

    @Override
    public void close() {
        if (client != null) {
            this.client.close();
        }
    }

    @Override
    public void stop() {
        this.close();
    }

    public static Builder builder() {
        return new Builder();
    }

    public static class Builder {
        private String uri;
        private String host;
        private Integer port;
        private String user;
        private String password;
        private String indexName;
        private String prefix;
        private Integer dimension;
        private Map<String, SchemaField> metadataConfig = new HashMap<>();

        public Builder uri(String uri) {
            this.uri = uri;
            return this;
        }

        public Builder host(String host) {
            this.host = host;
            return this;
        }

        public Builder port(Integer port) {
            this.port = port;
            return this;
        }

        public Builder user(String user) {
            this.user = user;
            return this;
        }

        public Builder password(String password) {
            this.password = password;
            return this;
        }

        public Builder indexName(String indexName) {
            this.indexName = indexName;
            return this;
        }

        public Builder prefix(String prefix) {
            this.prefix = prefix;
            return this;
        }

        public Builder dimension(Integer dimension) {
            this.dimension = dimension;
            return this;
        }

        public Builder metadataConfig(Map<String, SchemaField> metadataConfig) {
            this.metadataConfig = metadataConfig;
            return this;
        }

        public RedisEmbeddingStore build() {
            return this.uri != null ? new RedisEmbeddingStore(this.uri, this.indexName, this.prefix, this.dimension, this.metadataConfig)
                    : new RedisEmbeddingStore(this.host, this.port, this.user, this.password, this.indexName, this.prefix, this.dimension, this.metadataConfig);

        }
    }
}
